package org.canto.poss.instr.core.util;

import com.google.common.collect.ImmutableMap;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.RandomUtils;
import org.canto.poss.instr.core.abnormal.GeneralException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.dao.DuplicateKeyException;
import org.springframework.lang.NonNull;

import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;

/**
 * Description: 异步任务处理
 *
 * @author jsryin
 * @date 2024/12/13 7:34
 */
public class TaskExecutor {

    private static final Logger LOGGER = LoggerFactory.getLogger(TaskExecutor.class);
    public static final String BASE_EXECUTOR_NAME = "BaseExecutor";
    public static final String MONITOR_EXECUTOR_NAME = "default-inner-monitor-executor";
    private static final ScheduledExecutorService MONITOR_EXECUTOR = buildMonitorExecutor();
    private static final ThreadPoolExecutor BASE_EXECUTOR = buildThreadFirstExecutor(BASE_EXECUTOR_NAME);


    @NonNull
    private static ScheduledExecutorService buildMonitorExecutor() {
        ScheduledExecutorService executor = Executors.newScheduledThreadPool(getCpuProcessors());
        hookShutdownThreadPool(executor, MONITOR_EXECUTOR_NAME);
        return executor;
    }

    public static ThreadPoolExecutor getCommonExecutor() {
        return BASE_EXECUTOR;
    }

    public static ThreadPoolExecutor buildThreadFirstExecutor(String poolName) {
        int coreSize = getCpuProcessors();
        int maxSize = coreSize * 8;
        return buildThreadFirstExecutor(coreSize, maxSize, 5L, TimeUnit.MINUTES, 65536, poolName);
    }

    public static ThreadPoolExecutor buildThreadFirstExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, int workQueueSize, String poolName) {
        LinkedBlockingQueue<Runnable> queue = new LinkedBlockingQueue<Runnable>(workQueueSize) {
            private static final long serialVersionUID = 1L;

            public boolean offer(@NonNull Runnable o) {
                return false;
            }
        };
        RejectedExecutionHandler rejectedExecutionHandler = (runnable, executorx) -> {
            try {
                queue.put(runnable);
            } catch (InterruptedException var5) {
                LOGGER.warn("{} Queue offer interrupted. ", poolName, var5);
                Thread.currentThread().interrupt();
            }

        };
        ThreadPoolExecutor executor = new ThreadPoolExecutor(corePoolSize, maximumPoolSize, keepAliveTime, unit, queue, (new ThreadFactoryBuilder()).setNameFormat(poolName + "-%d").setUncaughtExceptionHandler((thread, throwable) -> {
            LOGGER.error("{} catching the uncaught exception, ThreadName: [{}]", new Object[]{poolName, thread.toString(), throwable});
        }).build(), rejectedExecutionHandler);
        executor.allowCoreThreadTimeOut(true);
        TaskExecutor.ExecutorManager.registerAndMonitorThreadPoolExecutor(poolName, executor);
        return executor;
    }

    public static <T> List<T> batchExecuteAsync(List<AsyncTask<T>> tasks, @NonNull String taskName) {
        return batchExecuteAsync(tasks, BASE_EXECUTOR, taskName);
    }

    public static <T> List<T> batchExecuteAsync(@NonNull List<AsyncTask<T>> tasks, @NonNull ThreadPoolExecutor executor, @NonNull String taskName) {
        if (CollectionUtils.isEmpty(tasks)) {
            return Collections.emptyList();
        } else {
            int size = tasks.size();
            List<Callable<T>> callables = tasks.stream().map((t) ->
                    (Callable<T>) () -> {
                        try {
                            T r = t.doExecute();
                            LOGGER.debug("[>>Executor<<] Async task execute success. ThreadName: [{}], BatchTaskName: [{}], SubTaskName: [{}]", new Object[]{Thread.currentThread().getName(), taskName, t.taskName()});
                            return r;
                        } catch (Throwable var3) {
                            LOGGER.warn("[>>Executor<<] Async task execute error. ThreadName: [{}], BatchTaskName: [{}], SubTaskName: [{}], exception: {}", new Object[]{Thread.currentThread().getName(), taskName, t.taskName(), var3.getMessage()});
                            throw var3;
                        }
                    }
            ).collect(Collectors.toList());
            CompletionService<T> cs = new ExecutorCompletionService(executor, new LinkedBlockingQueue(size));
            List<Future<T>> futures = new ArrayList(size);
            LOGGER.info("[>>Executor<<] Start async tasks, BatchTaskName: [{}], TaskSize: [{}]", taskName, size);
            Iterator var7 = callables.iterator();

            while (var7.hasNext()) {
                Callable<T> task = (Callable) var7.next();
                futures.add(cs.submit(task));
            }

            List<T> resultList = new ArrayList(size);

            for (int i = 0; i < size; ++i) {
                try {
                    Future<T> future = cs.poll(6L, TimeUnit.MINUTES);
                    if (future == null) {
                        cancelTask(futures);
                        LOGGER.error("[>>Executor<<] Async task [{}] - [{}] execute timeout, then cancel other tasks.", taskName, i);
                        throw new GeneralException("error.timeout");
                    }

                    T result = future.get();
                    resultList.add(result);
                    Object log;
                    if (result instanceof Collection) {
                        log = ((Collection) result).size();
                    } else {
                        log = result;
                    }

                    LOGGER.debug("[>>Executor<<] Async task [{}] - [{}] execute success, result: {}", new Object[]{taskName, i, log == null ? "null" : log});
                } catch (ExecutionException var12) {
                    LOGGER.warn("[>>Executor<<] Async task [{}] - [{}] execute error, then cancel other tasks.", new Object[]{taskName, i, var12});
                    cancelTask(futures);
                    Throwable throwable = var12.getCause();
                    if (throwable instanceof GeneralException) {
                        throw (GeneralException) throwable;
                    }

                    if (throwable instanceof DuplicateKeyException) {
                        throw (DuplicateKeyException) throwable;
                    }

                    throw new GeneralException("error.executorError", new Object[]{var12.getCause().getMessage()});
                } catch (InterruptedException var13) {
                    cancelTask(futures);
                    Thread.currentThread().interrupt();
                    LOGGER.error("[>>Executor<<] Async task [{}] - [{}] were interrupted.", taskName, i);
                    throw new GeneralException("error.error", new Object[0]);
                }
            }

            LOGGER.info("[>>Executor<<] Finish async tasks , BatchTaskName: [{}], TaskSize: [{}]", taskName, size);
            return resultList;
        }
    }

    public static void displayThreadPoolStatus(ThreadPoolExecutor threadPool, String threadPoolName) {
        displayThreadPoolStatus(threadPool, threadPoolName, (long) RandomUtils.nextInt(60, 600), TimeUnit.SECONDS);
    }

    public static void displayThreadPoolStatus(ThreadPoolExecutor threadPool, String threadPoolName, long period, TimeUnit unit) {
        TaskExecutor.CancelScheduledFuture cancelScheduledFuture = new TaskExecutor.CancelScheduledFuture();
        ScheduledFuture<?> scheduledFuture = MONITOR_EXECUTOR.scheduleAtFixedRate(() -> {
            String payload = "[>>ExecutorStatus<<] ThreadPool Name: [{}], Pool Status: [shutdown={}, Terminated={}], Pool Thread Size: {}, Largest Pool Size: {}, Active Thread Count: {}, Task Count: {}, Tasks Completed: {}, Tasks in Queue: {}";
            Object[] params = new Object[]{threadPoolName, threadPool.isShutdown(), threadPool.isTerminated(), threadPool.getPoolSize(), threadPool.getLargestPoolSize(), threadPool.getActiveCount(), threadPool.getTaskCount(), threadPool.getCompletedTaskCount(), threadPool.getQueue().size()};
            if (threadPool.getQueue().remainingCapacity() < 64) {
                LOGGER.warn(payload, params);
            } else {
                LOGGER.info(payload, params);
            }

            if (threadPool.isShutdown()) {
                cancelScheduledFuture.cancel();
            }

        }, 0L, period, unit);
        cancelScheduledFuture.setScheduledFuture(scheduledFuture);
    }

    public static void hookShutdownThreadPool(ExecutorService executorService, String threadPoolName) {
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            LOGGER.info("[>>ExecutorShutdown<<] Start to shutdown the thead pool: [{}]", threadPoolName);
            executorService.shutdown();

            try {
                if (!executorService.awaitTermination(60L, TimeUnit.SECONDS)) {
                    executorService.shutdownNow();
                    LOGGER.warn("[>>ExecutorShutdown<<] Interrupt the worker, which may cause some task inconsistent. Please check the biz logs.");
                    if (!executorService.awaitTermination(60L, TimeUnit.SECONDS)) {
                        LOGGER.error("[>>ExecutorShutdown<<] Thread pool can't be shutdown even with interrupting worker threads, which may cause some task inconsistent. Please check the biz logs.");
                    }
                }
            } catch (InterruptedException var3) {
                executorService.shutdownNow();
                LOGGER.error("[>>ExecutorShutdown<<] The current server thread is interrupted when it is trying to stop the worker threads. This may leave an inconsistent state. Please check the biz logs.");
                Thread.currentThread().interrupt();
            }

            LOGGER.info("[>>ExecutorShutdown<<] Finally shutdown the thead pool: [{}]", threadPoolName);
        }));
    }

    public static int getCpuProcessors() {
        return Runtime.getRuntime() != null && Runtime.getRuntime().availableProcessors() > 0 ? Runtime.getRuntime().availableProcessors() : 8;
    }

    private static <T> void cancelTask(List<Future<T>> futures) {
        Iterator iterator = futures.iterator();

        while (iterator.hasNext()) {
            Future<T> future = (Future) iterator.next();
            if (!future.isDone()) {
                future.cancel(true);
            }
        }

    }

    private static class CancelScheduledFuture {
        private volatile ScheduledFuture<?> scheduledFuture;

        private CancelScheduledFuture() {
        }

        public void setScheduledFuture(ScheduledFuture<?> scheduledFuture) {
            this.scheduledFuture = scheduledFuture;
        }

        public boolean cancel() {
            return Objects.isNull(this.scheduledFuture) ? false : this.scheduledFuture.cancel(true);
        }
    }

    public static class ExecutorManager {
        private static final ConcurrentHashMap<String, ThreadPoolExecutor> EXECUTORS = new ConcurrentHashMap(8);

        public ExecutorManager() {
        }

        public static void registerThreadPoolExecutor(String threadPoolName, ThreadPoolExecutor executor) {
            EXECUTORS.put(threadPoolName, executor);
        }

        public static void registerAndMonitorThreadPoolExecutor(String threadPoolName, ThreadPoolExecutor executor) {
            EXECUTORS.put(threadPoolName, executor);
            TaskExecutor.displayThreadPoolStatus(executor, threadPoolName);
            TaskExecutor.hookShutdownThreadPool(executor, threadPoolName);
        }

        public static ThreadPoolExecutor getThreadPoolExecutor(String threadPoolName) {
            return EXECUTORS.get(threadPoolName);
        }

        public static Map<String, ThreadPoolExecutor> getAllThreadPoolExecutor() {
            return ImmutableMap.copyOf(EXECUTORS);
        }

        public static void removeThreadPoolExecutor(String threadPoolName) {
            EXECUTORS.remove(threadPoolName);
        }
    }

}
